package com.itheima.stargraph.service.impl;

import cn.hutool.json.JSONUtil;
import com.itheima.stargraph.pojo.ComfyuiTask;
import com.itheima.stargraph.pojo.MessageBase;
import com.itheima.stargraph.service.ComfyuiMessageService;
import com.itheima.stargraph.service.ISgUserResultService;
import com.itheima.stargraph.service.Text2ImageService;
import com.itheima.stargraph.service.WsNotifyService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * Comfyui消息处理的业务层的实现类
 *
 * @Author mr.zhang
 * Date 2025/6/12 17:43
 */
@Service
public class ComfyuiMessageServiceImpl implements ComfyuiMessageService {

    /**
     * 类型：生图进度
     */
    private static final String TYPE_PROGRESS = "progress";

    /**
     * 类型：生图结果
     */
    private static final String TYPE_EXECUTED = "executed";

    /**
     * 类型：生图错误
     */
    private static final String TYPE_EXECUTION_ERROR = "execution_error";


    //注入推送消息的实例
    @Autowired
    private WsNotifyService wsNotifyService;
    //注入文生图实例
    @Autowired
    private Text2ImageService text2ImageService;



    /**
     * 根据消息的类型不同,推送消息给前端
     *
     * @param messageBase
     */
    @Override
    public void handleMsg(MessageBase messageBase) {
        //获取消息类型
        String type = messageBase.getType();
        //根据type来判断推送消息是什么类型
        if (type.equals(TYPE_PROGRESS)){
            this.handleProgress(messageBase);
        } else if (type.equals(TYPE_EXECUTED)) {
            this.handleExecured(messageBase);
        } else if (type.equals(TYPE_EXECUTION_ERROR)) {
            this.handleExecuredError(messageBase);
        }
    }

    /**
     * 推送生图错误
     * @param messageBase
     */
    private void handleExecuredError(MessageBase messageBase) {
        //封装数据
        HashMap<String, Object> map = new HashMap<>();
        map.put("type",TYPE_EXECUTION_ERROR);

        Object promptIdObj = map.getOrDefault("prompt_id", null);
        if (promptIdObj==null){
            return;
        }

        String promptId = String.valueOf(promptIdObj);

        ComfyuiTask comfyuiTask = text2ImageService.getStartedTask(promptId);
        String clientId = comfyuiTask.getWsClientId();//前端的clientId

        wsNotifyService.sendMsgToUser(clientId,JSONUtil.toJsonStr(map));


    }


    @Autowired
    private ISgUserResultService iSgUserResultService;

    /**
     * 推送生图结果
     * @param messageBase
     */
    private void handleExecured(MessageBase messageBase) {
        //获取data数据
        HashMap<String, Object> data = messageBase.getData();
        //判断是否有prompt_id
        Object promptIdObj = data.getOrDefault("prompt_id", null);
        if (promptIdObj==null){
            return;
        }
        String promptId = String.valueOf(promptIdObj);
        //获取缓存中最开始的任务
        ComfyuiTask comfyuiTask = text2ImageService.getStartedTask(promptId);
        if (comfyuiTask==null){
            return;
        }
        String clientId = comfyuiTask.getWsClientId();//获取客户端id,就是前端的clientId
        //组装数据
        data.put("type","imageResult");
        Object output = data.get("output");
        Map<String, Object> outputMap = (Map<String, Object>) output;

        Object images = outputMap.get("images");
        List<Map<String, Object>> imageList = (List<Map<String, Object>>) images;

        List<String> stringList = imageList.stream().map(x -> {
            return String.format("http://localhost:8188/view?filename=%s&type=%s&&subfolder=%s", x.get("filename"), x.get("type"), x.get("subfolder"));
        }).collect(Collectors.toList());
        data.put("urls",stringList);


        //推送消息给前端
        wsNotifyService.sendMsgToUser(clientId,JSONUtil.toJsonStr(data));

        //将生图结果存入数据库中
        iSgUserResultService.insertExecured(comfyuiTask.getUserId(),stringList);
    }

    /**
     * 推送生图进度
     * @param messageBase
     */
    private void handleProgress(MessageBase messageBase) {
        //获取messageBase中数据
        HashMap<String, Object> data = messageBase.getData();
        //将type作为键,添加到获取的数据中,获得一个大的json对象
        data.put("type",TYPE_PROGRESS);

        //根据缓存的promptId获取clientId
        Object promptIdObj = data.getOrDefault("prompt_id",null);
        if (promptIdObj==null){
            return;
        }
        String promptId = String.valueOf(promptIdObj);
        ComfyuiTask comfyuiTask = text2ImageService.getStartedTask(promptId);
        String clientId = comfyuiTask.getWsClientId();//获取到的前端的clientId

        //推送消息给前端
        wsNotifyService.sendToUser(clientId, JSONUtil.toJsonStr(data));
    }
}
